// server/main.go
package main

import (
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源，仅用于测试
	},
}

// 用于存储所有活跃的 WebSocket 连接
var (
	clients     = make(map[*websocket.Conn]bool)
	clientsLock sync.RWMutex
)

// 消息广播器
func broadcaster() {
	ticker := time.NewTicker(10 * time.Millisecond) // 10ms 发送一次，每秒 100 次
	defer ticker.Stop()

	counter := 0
	for range ticker.C {
		counter++
		msg := fmt.Sprintf("消息 #%d 生成于 %s", counter, time.Now().Format("15:04:05.000"))

		clientsLock.RLock()
		// 遍历所有客户端并发送消息
		for conn := range clients {
			// 使用 goroutine 避免单个慢客户端阻塞整个广播
			go func(c *websocket.Conn) {
				if err := c.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
					log.Printf("向客户端发送消息失败: %v", err)
					// 标记连接需要关闭
					clientsLock.Lock()
					delete(clients, c)
					clientsLock.Unlock()
					c.Close()
				}
			}(conn)
		}
		clientsLock.RUnlock()
	}
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	// 升级 HTTP 连接到 WebSocket
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return
	}
	defer conn.Close()

	// 将新连接添加到客户端列表
	clientsLock.Lock()
	clients[conn] = true
	clientCount := len(clients)
	clientsLock.Unlock()

	log.Printf("新客户端连接。当前连接数: %d", clientCount)

	// 注意：服务端不读取客户端消息，只负责发送
	// 但需要持续读取以响应 ping/pong 和检测连接状态
	// 否则客户端的 ping 会积压导致连接被关闭
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			// 连接可能已关闭
			log.Printf("读取客户端消息时出错 (可能是正常关闭): %v", err)
			break
		}
		// 这里可以处理客户端发来的消息，但本例不需要
	}

	// 连接断开，从列表中移除
	clientsLock.Lock()
	delete(clients, conn)
	clientCount = len(clients)
	clientsLock.Unlock()
	log.Printf("客户端断开连接。剩余连接数: %d", clientCount)
}

func main() {
	// 启动广播器 goroutine
	go broadcaster()

	http.HandleFunc("/ws", wsHandler)

	log.Println("WebSocket 服务端启动，监听 :8080")
	log.Println("客户端连接地址: ws://localhost:8080/ws")

	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal("服务器启动失败: ", err)
	}
}
